#ifndef _SPATIAL_TREE_H_
#define _SPATIAL_TREE_H_

#include "../cuCommon/cuTypes.h"
#include "../cuCommon/cuBuffer.h"

//#include "../glCommon/glRender.h"

#include "../rtCommon/Histopyramid.h"
#include "../rtCommon/Axis.h"
#include "../rtCommon/Box.h"
#include "../rtCommon/Benchmark.h"

#include "../xobjIO/glWorld.cuh"

#define TREE_AXIS_X					   0x00
#define TREE_AXIS_Y					   0x01
#define TREE_AXIS_Z					   0x02

#define TREE_LEAF              0x03
#define TREE_NON_LEAF          0x00

#define TREE_ROOT              0x00

//////////////////////////////////////////////////////////////////////////
/// SInterval structure 
//////////////////////////////////////////////////////////////////////////
struct SInterval
{
	float start, end;

	SInterval() { start = 0; end = 0; }
	SInterval(float l, float r) { start = l; end = r; }

	SInterval & operator = (const float2     &i) { start = i.x;     end = i.y;   return *this; }
	SInterval & operator = (const SInterval &i) { start = i.start; end = i.end; return *this; }

	operator float2() { return f2(start, end); }

	float length() { return end - start; }
};

//////////////////////////////////////////////////////////////////////////
/// Tree Node structure used during build stage (aligned to 128 bits)
//////////////////////////////////////////////////////////////////////////
struct TreeNode
{
	//////////////////////////////////////////////////////////////////////////
	/// split position with axis coded in lower bits
	//////////////////////////////////////////////////////////////////////////
	union
	{
		float   split; // store splitting position
		uint32  bin;   // or bin index
	};

	//////////////////////////////////////////////////////////////////////////
	/// Start Index / Left Node union
	//////////////////////////////////////////////////////////////////////////
	union 
	{
		uint32  start;  
		uint32  left;
	};

	//////////////////////////////////////////////////////////////////////////
	/// Start Index 
	//////////////////////////////////////////////////////////////////////////
	union
	{
		uint32  end;  
		uint32  right;
	};

	//////////////////////////////////////////////////////////////////////////
	/// parent index to build traversal tree
	//////////////////////////////////////////////////////////////////////////
	union 
	{
		uint32 parent;
		float  sah;
	};

	TreeNode()
	{
		split = as_float(TREE_NON_LEAF);
		start = 0;
		end   = 0;
		parent = TREE_ROOT;
	}

	TreeNode(float s)
	{
		split = s;
		start = 0;
		end   = 0;
		parent = TREE_ROOT;
	}

	TreeNode(uint32 p)
	{
		split = as_float(TREE_NON_LEAF);
		start = 0;
		end   = 0;
		parent = p;
	}

	TreeNode(float s, uint32 p)
	{
		split = s;
		start = 0;
		end   = 0;
		parent = p;
	}

	TreeNode(uint32 s, uint32 e)
	{
		split = as_float(TREE_NON_LEAF);
		start = s;
		end   = e;
		parent = TREE_ROOT;
	}

	bool Leaf()
	{
		return (as_uint(split) & TREE_LEAF) == TREE_LEAF;
	}

	void SplitAxis(float s, int axis)
	{
		split = as_float(as_uint(s) & (~TREE_LEAF) | axis);
	}

	float Split()
	{
		return as_float(as_uint(split) & (~TREE_LEAF));
	}

	uint Axis()
	{
		return         (as_uint(split) &  TREE_LEAF );
	}
};

struct KdNode
{
	union 
	{
		float    split;      // for non leaf nodes TN holds split position with axis encoded in lower bits
		uint32   count;      // for leaf node there is no split, but there is primitive count stores shifted right by two
		float    t;          // for tracing t stores parametric intersection data.
	};

	union 
	{
		uint32   right;      // for non leaf nodes TN holds offset to the right sub-node
		uint32   start;      // for leaf node there is no right sub-node, but there is an index into primIdx soup
		uint32   next;       // for tracing next stores stacked index
	};

	KdNode()
	{
		split = 0;
		right = 0;
	}
	
	void SplitAxis(float s, int axis)
	{
		split = as_float(as_uint(s) & (~TREE_LEAF) | axis);
	}

	float Split()
	{
		return as_float(as_uint(split) & (~TREE_LEAF));
	}

	uint Axis()
	{
		return (count & TREE_LEAF);
	}

	uint Count()
	{
		return (count >> 2);
	}
};

struct BVHNode
{
	union 
	{
		struct
		{
			float3 pmin;
			union
			{
				uint32 right;
				uint32 start;
			};
		};
		float4 pmin_rs;
	};

	union
	{
		struct
		{
			float3 pmax;
			union
			{
				float  split;
				uint32 axis;
				uint32 count;
			};
		};
		float4 pmax_lc;
	};

	void SplitAxis(float s, int axis)
	{
		split = as_float(as_uint(s) & (~TREE_LEAF) | axis);
	}

	uint Axis()
	{
		return (axis & TREE_LEAF );
	}

	uint Count()
	{
		return (count >> 2);
	}
	
	float Split()
	{
		return as_float(as_uint(split) & (~TREE_LEAF));
	}
};

class CSpatialTree
{
	friend class CglWorld;

protected:
	
	sf4Buffer     * m_pPos;                         // pointer to position packed in f4
	si_Buffer     * m_pPosIdx;                      // pointer to regular position index
	ui_Buffer     * m_pPosIdxSwzld;                 // pointer to swizzled position index
	uint32          m_nFace;                        // number of faces in the scene
	
	sf2Buffer       m_PrimitiveBI[XYZ];             // per primitive bounding interval values packed in 3 different axis arrays

	f2HistoPyramid  m_histopyramid[XYZ];            // l-r histo-pyramid 

	ui_Buffer       m_primIdx;                      // large index buffer that is used for tree construction - stores intermidiate data
	uint32          m_nPrimitive;                   // m_nPrimitive is used as an index into m_primIdx arrays
	ui_Buffer       m_treeIdx;                      // large index buffer that is used for tree construction - stores results
	uint32          m_nTree;                        // m_nTree is used as an index into m_treeIdx

	sf4Buffer       m_Tree;                         // m_Tree is an internal data structure that couples with m_Interval to construct 
	sf2Buffer       m_Interval;                     // various data structures such as KdTree or BVH

	sf2Buffer       m_KdTree;                       // KdTree is a packed array of f2 that can be cast to KdNode
	sf4Buffer       m_BVHTree;                      // BVH    is a packed array of f8 that can be cast to BVHNode
 
	S_Box          m_box;                          // Bounding box for the whole scene

	int32           m_nMaxDepth;
	int32           m_nMaxPrim;                     // Maximum primitives inside a node
	
	// m_nPrimExplode value determines how big the prealocated arrays 
	// m_primIdx and m_treeIdx should be relative to the whole scene
	uint32          m_nPrimExplode;                 

	int32           m_nNodeIndex;
	int32           m_nInterval;
	int32           m_nTraversalIdx;

	void            Null();

	float           EvaluateSurfaceAreaHeuristics(float probabilityLeft, float probabilityRight, int32 triangleCountLeft, int32 triangleCountRight, float Ct = 1.0f, float Cb = 1.0f);
	float           EvaluateMedianHeuristics(int32 triangleCountLeft, int32 triangleCountRight, float Ct = 1.0f, float Cb = 1.0f);

public:

	CSpatialTree();
	~CSpatialTree();

	//////////////////////////////////////////////////////////////////////////
	/// Create function allocates necessary resources
	//////////////////////////////////////////////////////////////////////////
	void            Create(sf4Buffer * pPos, ui_Buffer * pPosIdxSwzld, si_Buffer * pPosIdx);

	//////////////////////////////////////////////////////////////////////////
	/// Release resources
	//////////////////////////////////////////////////////////////////////////
	void            Release();

	//////////////////////////////////////////////////////////////////////////
	/// 1. Compute per primitive AABB and store them in split[XYZ]
	//////////////////////////////////////////////////////////////////////////
	void            PrimitiveAABB(uint8 exeFlags);

	//////////////////////////////////////////////////////////////////////////
	/// 3. Compute Min Max Histogram Pyramid and Scan for KD Tree binning
	//////////////////////////////////////////////////////////////////////////
	void            Histogram(uint8        exeFlags,
														uint32     * pPrim, 
														uint32       nPrim,
														SInterval * pInterval);

	//////////////////////////////////////////////////////////////////////////
	/// 4. Evaluate Surface Area Heuristics
	//////////////////////////////////////////////////////////////////////////
	bool            SAH(uint8 exeFlags, SInterval * pInterval, TreeNode & s, uint32 & axis, int32 & sah_bin, uint32 nPrim);

	//////////////////////////////////////////////////////////////////////////
	/// KdTree_p0 function creates KdTree Tree data structure using SAH
	//////////////////////////////////////////////////////////////////////////
	void            KdTree_p0(uint8 exeFlags);

	//////////////////////////////////////////////////////////////////////////
	/// KdTree_p1 creates a KdTree data structure from Tree
	//////////////////////////////////////////////////////////////////////////
	void            KdTree_p1(uint8 exeFlags);

	//////////////////////////////////////////////////////////////////////////
	/// KdTree_p2 creates a KdTree data structure from Tree recursively
	//////////////////////////////////////////////////////////////////////////
	void            KdTree_p2(TreeNode *pb, KdNode * pt, TreeNode & n);

	//////////////////////////////////////////////////////////////////////////
	/// BVHTree_p0 function creates Tree data structure using SAH
	//////////////////////////////////////////////////////////////////////////
	void            BVHTree_p0(uint8 exeFlags);

	//////////////////////////////////////////////////////////////////////////
	/// BVHTree_p1 creates a BVH data structure from Tree
	//////////////////////////////////////////////////////////////////////////
	void            BVHTree_p1(uint8 exeFlags);

	//////////////////////////////////////////////////////////////////////////
	/// BVHTree_p0 creates a BVH data structure from Tree recursively
	//////////////////////////////////////////////////////////////////////////
	void            BVHTree_p2(TreeNode *pb, BVHNode * pt, SInterval * pi, TreeNode & n, int i);

	//////////////////////////////////////////////////////////////////////////
	/// Get TraversalData
	//////////////////////////////////////////////////////////////////////////
	float2 *        pKdTree();
	float4 *        pBVHTree();

	uint   *        pPrimitive();
	int    *        pFace();
	float4 *        pVertex();

	//////////////////////////////////////////////////////////////////////////
	/// Get MaxDepth
	//////////////////////////////////////////////////////////////////////////
	int32           maxNode();

	//////////////////////////////////////////////////////////////////////////
	/// Get SAABox
	//////////////////////////////////////////////////////////////////////////
	S_Box           box();

	//////////////////////////////////////////////////////////////////////////
	/// Render methods
	//////////////////////////////////////////////////////////////////////////
	//void RenderKdTree(int maxLevel);
};

#endif _SPATIAL_TREE_H_